from volcenginesdkarkruntime import Ark

class Doubao:
    def __init__(self, base_url, api_key, initial_prompt):
        """
        类的初始化方法，用于创建 Ark 客户端实例并设置初始消息列表。
        :param base_url: Ark 服务的基础 URL
        :param api_key: 调用 Ark 服务所需的 API 密钥
        :param initial_prompt: 初始的系统提示词，用于设置对话的初始上下文
        """
        self.client = Ark(
            base_url=base_url,
            api_key=api_key
        )
        self.initial_prompt = initial_prompt
        self.estimated_length = 400
        self.messages = [
            {"role": "system", "content": initial_prompt}
        ]

    def run(self, user_prompt):
        """
        向模型发送用户提示词并以信息流形式返回回复内容。
        :param user_prompt: 用户输入的提示词，用于向模型提问
        :return: 模型生成的回复内容的信息流
        """
        self.messages.append({"role": "user", "content": user_prompt})
        stream = self._initial_text_generation_stream(self.messages)
        text = ""
        for chunk in stream:
            if chunk:
                text += chunk
                yield chunk
        self.messages.append({"role": "assistant", "content": text})

    def generate_text(self, user_prompt):
        """
        向模型发送用户提示词并生成回复文本，同时会保留对话历史。
        :param user_prompt: 用户输入的提示词，用于向模型提问
        :return: 模型生成的回复文本
        """
        self.messages.append({"role": "user", "content": user_prompt})
        text = self._initial_text_generation(self.messages)
        return text

    def generate_text_arr(self, user_prompt):
        """
        向模型发送用户提示词并生成回复文本，同时会保留对话历史。
        :param user_prompt: 用户输入的提示词，用于向模型提问
        :return: 模型生成的回复文本
        """
        ret=[]
        for item in user_prompt:

            self.messages.append({"role": "user", "content": user_prompt})
            text = self._initial_text_generation(item)
            ret.append(text)
        return ret

    def _initial_text_generation_stream(self, data):
        """
        实际调用 Ark 客户端的 API 来生成文本的内部方法，以流的形式返回。
        :param data: 包含对话历史的消息列表，用于传递给模型
        :return: 模型生成的回复内容的信息流
        """
        stream = self.client.chat.completions.create(
            model="ep-20250216183741-sncg4",
            messages=data,
            stream=True
        )
        for chunk in stream:
            if not chunk.choices:
                continue
            delta_content = chunk.choices[0].delta.content
            if delta_content:
                yield delta_content

    def _initial_text_generation(self, data):
        """
        实际调用 Ark 客户端的 API 来生成文本的内部方法。
        :param data: 包含对话历史的消息列表，用于传递给模型
        :return: 模型生成的完整回复文本
        """
        stream = self._initial_text_generation_stream(data)
        text = ""
        for chunk in stream:
            text += chunk
        return text

    def reset_and_reask(self, new_user_prompt):
        """
        清空之前的对话历史，仅保留初始系统提示词，然后使用新的提示词重新发起询问。
        :param new_user_prompt: 用于重新提问的新提示词
        :return: 模型针对新提示词生成的回复文本
        """
        # 清空历史消息，只保留初始系统提示词
        self.messages = [
            {"role": "system", "content": self.initial_prompt}
        ]
        return self.generate_text(new_user_prompt)


if __name__ == "__main__":
    initial_prompt = """深厚积累了有关 黄历玄学知识、 短视频脚本创作时，你需要充分融合这些专业背景知识； 根据用户输入的主题创作需求，进行短视频脚本创作  多多使用口语容易让人理解的话  开头为 "日日看黄历 事事皆顺利"  需要详细的 以下 值日 做什么类型的事情比较顺利 什么是不建议做(结合生活) 值神是谁 这一天需要注意的是 建议做或不建议做的事情 (结合生活)   说一下 属相相冲的  根据幸运色推荐衣服 和麻将的方位注意避免的 不要带 "儿" 这个字 根据输入的消息来润色脚本不需要写镜头  画面 和天气 也不要输出表格 文章 410 字左右"""
    api = "https://ark.cn-beijing.volces.com/api/v3"
    api_key = "c837ff21-6e5f-49ba-960e-66e2edfe42c7"

    doubao = Doubao(api, api_key, initial_prompt)
    # 以字符串形式获取回复
    # mstr = doubao.generate_text("今天是农历九月十四")
    # print(mstr)

    # 以信息流形式获取回复
    print("以下是信息流形式的回复：")
    for chunk in doubao.run("今天是农历九月十四"):
        print(chunk, end="", flush=True)